
import pandas as pd
import numpy as np

import os
import json
import re
from tqdm import tqdm

from transformers import *
import tensorflow as tf
import pandas as pd
import numpy as np
from tqdm import tqdm
import tensorflow_addons as tfa
import os
import json
import re
from sklearn.model_selection import KFold

from model import build_model
from data_process import *


data_prefix = "/data2/workspace/2021_comp/ccks-ner/data/base/"

label_dic_path = data_prefix+"label_dic.json"
combain_label_dic_path = data_prefix+"combain_label_dic.json"
test_path = data_prefix+"final_test.txt"
train_path = data_prefix+"train.conll"
val_path = data_prefix+"dev.conll"
label_dic = json.load(open(label_dic_path))
reverse_label_dic = dict([(item[1], item[0]) for item in label_dic.items()])

def write_k_fold_data(fold_path: str, kfold_num: int):
    """Writes the kfold data to disk .

    Args:
        fold_path (str): [description]
        kfold_num (int): [description]
    """

    kfold = KFold(kfold_num)

    seq_lists, entity_lists = get_seq_entity(
        open(train_path).readlines()+open(val_path).readlines())

    for fold_id, (train_idxs, val_idxs) in enumerate(kfold.split(seq_lists)):

        fold_k = os.path.join(fold_path, f"fold_{fold_id}")
        assert not os.path.exists(fold_k)

        os.makedirs(fold_k)

        train_file = os.path.join(fold_path, f"fold_{fold_id}", "train.conll")
        val_file = os.path.join(fold_path, f"fold_{fold_id}", "dev.conll")
        train_list = []
        val_list = []

        for idx in train_idxs:
            text = seq_lists[idx]
            entitys = entity_lists[idx]
            assert len(text) == len(entitys)
            for item in zip(text, entitys):
                train_list.append(" ".join(item)+"\n")

            train_list.append("\n")

        for idx in val_idxs:
            text = seq_lists[idx]
            entitys = entity_lists[idx]
            assert len(text) == len(entitys)
            for item in zip(text, entitys):
                val_list.append(" ".join(item)+"\n")

            val_list.append("\n")

        open(train_file, mode="w").writelines(train_list)

        open(val_file, mode="w").writelines(val_list)

        print(train_file, "len ", len(train_list))
        print(val_file, "len ", len(val_list))


# write_k_fold_data("./data/5fold",5)
# write_k_fold_data("./data/3fold",3)
# write_k_fold_data("./data/7fold",7)
# write_k_fold_data("./data/10fold",10)